import torchvision.datasets
from torchvision import transforms
import torch
from torch.utils.data import DataLoader
import torch.nn as nn
import torchvision

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
ROOT_DIR = "/Users/leslie/PycharmProjects/paper_audio/data"
ROOT_DIR_TRAIN = "/home/ubuntu/leslie/data_train"
ROOT_DIR_TEST = "/home/ubuntu/leslie/data_test"




class AlexNet(nn.Module):
    def __init__(self, num_classes=6):
        super(AlexNet, self).__init__()
        self.feature_extraction = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=96, kernel_size=11, stride=4, padding=2, bias=False),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=0),
            nn.Conv2d(in_channels=96, out_channels=192, kernel_size=5, stride=1, padding=2, bias=False),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=0),
            nn.Conv2d(in_channels=192, out_channels=384, kernel_size=3, stride=1, padding=1, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, stride=1, padding=1, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1, bias=False),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=0),
        )
        self.classifier = nn.Sequential(
            nn.Dropout(p=0.5),
            nn.Linear(in_features=256 * 6 * 6, out_features=4096),
            nn.Dropout(p=0.5),
            nn.Linear(in_features=4096, out_features=4096),
            nn.Linear(in_features=4096, out_features=num_classes),
            nn.Softmax()
        )

    def forward(self, x):
        x = self.feature_extraction(x)
        x = x.view(x.size(0), 256 * 6 * 6)
        x = self.classifier(x)
        return x


def data_prepare():
    # target_transform = transforms.Lambda(lambda y: torch.zeros(10, dtype=torch.float).scatter_(dim=0, index=torch.tensor(y), value=1))

    data_transforms = transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    training_data = torchvision.datasets.ImageFolder(ROOT_DIR,
                                                     transform=data_transforms,
                                                     target_transform=None)

    test_data = torchvision.datasets.ImageFolder(ROOT_DIR,
                                                 transform=data_transforms,
                                                 target_transform=None)
    return training_data, test_data


def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):



        # Compute prediction and loss
        X = X.to(device)
        pred = model(X)
        y = y.to(device)
        pred = pred.to(device)
        loss = loss_fn(pred, y)




        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")




def test_loop(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:




            X = X.to(device)
            y = y.to(device)
            pred = model(X).to(device)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")


def train():
    training_data, test_data = data_prepare()
    train_dataloader = DataLoader(training_data, batch_size=2, shuffle=True)
    test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)

    # model = torchvision.models.AlexNet()
    model = AlexNet()

    model.to(device)
    # print(model)

    # Initialize
    learning_rate = 1e-3
    batch_size = 10
    epochs = 200
    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

    epochs = 100
    for t in range(epochs):
        print(f"Epoch {t + 1}\n-------------------------------")
        train_loop(dataloader=train_dataloader, model=model, loss_fn=loss_fn, optimizer=optimizer)
        test_loop(dataloader=test_dataloader, model=model, loss_fn=loss_fn)
    print("Done!")

if __name__ == '__main__':
    train()






